package vip.liux.backend.infrastructure.utils;

import com.github.benmanes.caffeine.cache.Cache;
import org.slf4j.Logger;
import org.springframework.util.ResourceUtils;
import vip.liux.cache.GlobalCache;

import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.security.*;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

import static org.slf4j.LoggerFactory.getLogger;

public class KeyGenerator {

    private static final String PUBLIC_KEY_PATH = getSavePath() + "public.cer";
    private static final String PRIVATE_KEY_PATH = getSavePath() + "private.key";

    private static final Logger logger = getLogger(KeyGenerator.class);
    private static final Cache<String, Object> cache = GlobalCache.getCache("KeyGenerator", Long.MAX_VALUE);

    public static void main(String[] args) {
        generateKeys();
    }

    // 获取本地私钥文件
    public static String getPrivateKeyString() {
        return (String) cache.get("privateKeyString", k -> {
            try {
                byte[] keyBytes = Files.readAllBytes(Paths.get(PRIVATE_KEY_PATH));
                String keyString = new String(keyBytes);
                return keyString
                        .replace("-----BEGIN PRIVATE KEY-----", "")
                        .replace("-----END PRIVATE KEY-----", "")
                        .replaceAll("\\s", "");
            } catch (IOException e) {
                logger.warn("Error reading private key file: {}", e.getMessage());
                generateKeys();
                return getPrivateKeyString();
            }
        });
    }

    public static PrivateKey getPrivateKey() {
        return (PrivateKey) cache.get("privateKey", k -> {
            try {
                PKCS8EncodedKeySpec spec = new PKCS8EncodedKeySpec(Base64.getDecoder().decode(getPrivateKeyString()));
                KeyFactory keyFactory = KeyFactory.getInstance("RSA");
                return keyFactory.generatePrivate(spec);
            } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
                throw new RuntimeException("Error getting private key: " + e.getMessage());
            }
        });
    }

    public static String getPublicKeyString() {
        return (String) cache.get("publicKeyString", k -> {
            try {
                byte[] keyBytes = Files.readAllBytes(Paths.get(PUBLIC_KEY_PATH));
                String keyString = new String(keyBytes);
                return keyString
                        .replace("-----BEGIN PUBLIC KEY-----", "")
                        .replace("-----END PUBLIC KEY-----", "")
                        .replaceAll("\\s", "");
            } catch (IOException e) {
                logger.warn("Error reading public key file: {}", e.getMessage());
                generateKeys();
                return getPublicKeyString();
            }
        });
    }

    public static PublicKey getPublicKey() {
        return (PublicKey) cache.get("publicKey", k -> {
            try {
                X509EncodedKeySpec spec = new X509EncodedKeySpec(Base64.getDecoder().decode(getPublicKeyString()));
                KeyFactory keyFactory = KeyFactory.getInstance("RSA");
                return keyFactory.generatePublic(spec);
            } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
                throw new RuntimeException("Error getting public key: " + e.getMessage());
            }
        });
    }

    private static void generateKeys() {
        try {
            // Generate the key pair
            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
            keyPairGenerator.initialize(2048);
            KeyPair keyPair = keyPairGenerator.generateKeyPair();

            // Get the public and private keys
            PublicKey publicKey = keyPair.getPublic();
            PrivateKey privateKey = keyPair.getPrivate();

            // Save the keys to the resources directory
            saveKeyToFile(PUBLIC_KEY_PATH, publicKey.getEncoded(), false);
            saveKeyToFile(PRIVATE_KEY_PATH, privateKey.getEncoded(), true);

            logger.info("Keys generated and saved to resources directory.");
        } catch (NoSuchAlgorithmException | IOException e) {
            logger.warn("Error generating keys: {}", e.getMessage());
        }
    }

    private static String getSavePath() {
        try {
            return ResourceUtils.getFile("classpath:").getParentFile().getParentFile() + "\\src\\main\\resources\\wwwroot\\keys\\";
        } catch (FileNotFoundException e) {
            if (logger != null) {
                logger.warn("Error getting save path: {}", e.getMessage());
            }
        }
        return "D:\\wwwroot\\keys\\";
    }

    private static void saveKeyToFile(String filePath, byte[] key, boolean isPrivateKey) throws IOException {
        Files.createDirectories(Paths.get(filePath).getParent());
        try (FileOutputStream fos = new FileOutputStream(filePath)) {
            String keyString = Base64.getEncoder().encodeToString(key);
            if (isPrivateKey) {
                keyString = "-----BEGIN PRIVATE KEY-----\n" + keyString + "\n-----END PRIVATE KEY-----";
            } else {
                keyString = "-----BEGIN PUBLIC KEY-----\n" + keyString + "\n-----END PUBLIC KEY-----";
            }
            fos.write(keyString.getBytes());
        }
    }
}